from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import fsolve

plt.rcParams['axes.unicode_minus'] = False
plt.rc('font', size=10)
plt.rc('font', family='SimHei')

v = 1
b = 0.55 / (2 * np.pi)
theta_0 = 16 * 2 * np.pi
l_head = 3.41 - 0.275 * 2
l_body = 2.2 - 0.275 * 2
l = 0.275
half_width = 0.15
label_list = [0, 1, 51, 101, 151, 201, 222]


# theta与t的关系
def pfun(y, t):
    global v, b
    return (-1) * v / (b * np.sqrt(1 + y ** 2))


def x(theta, r):
    return r * np.cos(theta)


def y(theta, r):
    return r * np.sin(theta)


# 计算速度向量的函数
def velocity_from_theta(theta, dtheta_dt, b):
    u_x = np.cos(theta) - theta * np.sin(theta)
    u_y = np.sin(theta) + theta * np.cos(theta)
    return b * dtheta_dt * np.array([u_x, u_y])


"""
求解极坐标系下圆和等距螺线的交点
"""
def find_intersections(r0, phi0, R, b, theta_guesses):
    def equations(p):
        theta = p[0]
        r = b * theta
        eq1 = r ** 2 + r0 ** 2 - 2 * r * r0 * np.cos(theta - phi0) - R ** 2
        return [eq1]

    solutions = []
    for guess in theta_guesses:
        theta_sol = fsolve(equations, guess)[0]
        r_sol = b * theta_sol
        if theta_sol > 0 and abs(equations([theta_sol])[0]) < 1e-6:
            is_new = True
            for sol in solutions:
                if np.allclose([theta_sol], [sol[0]], atol=1e-3):
                    is_new = False
                    break
            if is_new:
                solutions.append((theta_sol, r_sol))
    return solutions


# 修改后的find_next函数，现在返回位置和速度，并递归计算后续棍头
def find_next(prev_theta, prev_r, prev_pos_cart, prev_v, i=0):
    # 使用前一个棍头的极坐标作为圆心
    theta_guesses = [prev_theta + 0.1]  # 初始猜测角度
    if i == 0:
        intersections = find_intersections(prev_r, prev_theta, l_head, b, theta_guesses)
    else:
        intersections = find_intersections(prev_r, prev_theta, l_body, b, theta_guesses)

    results = []
    for theta_sol, r_sol in intersections:
        # 当前棍头的笛卡尔坐标
        pos_cart = np.array([x(theta_sol, r_sol), y(theta_sol, r_sol)])
        # 计算向量 delta_r = current - previous
        delta_r = pos_cart - prev_pos_cart

        # 计算当前速度
        current_v = prev_v * (prev_r * (np.cos(prev_theta) ** 2) - r_sol * np.cos(prev_theta) * np.cos(theta_sol) - prev_r * prev_theta * np.sin(prev_theta) * np.cos(prev_theta) + r_sol * prev_theta * np.sin(prev_theta) * np.cos(theta_sol)) / (prev_r * np.cos(prev_theta) * np.cos(theta_sol) - r_sol * (np.cos(theta_sol) * 2) - prev_r * theta_sol * np.cos(prev_theta) * np.sin(theta_sol) + r_sol * theta_sol * np.sin(theta_sol) * np.cos(theta_sol))

        # 存储结果：位置、速度、极角、极径
        results.append({
            'theta': theta_sol,
            'r': r_sol,
            'pos_cart': pos_cart,
            'velocity': current_v,
            'index': i
        })

        if i > 222 or theta_sol > soli[0]:
            return []
        else:
            # print(pos_cart)
            if i in label_list:
                print(f"i={i}时，位置：{x(prev_theta, prev_r), y(prev_theta, prev_r)}，速度：{prev_v}")
            # 递归计算下一个棍头，传递当前棍头的参数
            next_results = find_next(theta_sol, r_sol, pos_cart, current_v, i + 1)
            results.extend(next_results)
    return results


# 主程序
t = np.arange(0, 400, 0.1)
soli = odeint(pfun, theta_0, t)[:, 0]
# 选择时间点，例如t_index=600（对应t=60s，因为步长0.1，索引600）
t_index = 360
theta0 = soli[t_index * 10]
r0 = b * theta0
dtheta0_dt = pfun(theta0, t[t_index * 10])
pos0_cart = np.array([x(theta0, r0), y(theta0, r0)])

# 计算所有棍头的位置和速度
all_heads = find_next(theta0, r0, pos0_cart, v)

# 绘制图形
plt.figure(figsize=(8, 8))
# 绘制螺线
theta_plot = np.linspace(0, soli[0], 1000)
r_plot = b * theta_plot
plt.plot(x(theta_plot, r_plot), y(theta_plot, r_plot), label=f'螺线: r = {b}θ')
# 绘制第一个棍头
plt.plot(pos0_cart[0], pos0_cart[1], 'ro')
# 绘制其他棍头
for i in range(len(all_heads)):
    pos = all_heads[i]['pos_cart']
    plt.plot(pos[0], pos[1], 'ro')
    if i > 0:
        prev_pos = all_heads[i - 1]['pos_cart']
        plt.plot([prev_pos[0], pos[0]], [prev_pos[1], pos[1]], color='red')
    else:
        plt.plot([pos0_cart[0], pos[0]], [pos0_cart[1], pos[1]], color='red')
    # 绘制速度向量
    # v = head['velocity']
    # print(f"位置：{pos}，速度向量：{v}，速度大小：{(v[0] ** 2 + v[1] ** 2) ** 0.5}")
    # plt.arrow(pos[0], pos[1], v[0] * 0.1, v[1] * 0.1, head_width=0.1, head_length=0.2, fc='g', ec='g')

plt.xlabel('X')
plt.ylabel('Y')
plt.title('把手位置')
plt.axis('equal')
plt.grid(True)
plt.legend()
plt.show()
